import matplotlib as mpl
mpl.rcParams['font.sans-serif'] = ['SimHei']
mpl.rcParams['axes.unicode_minus'] = False
import pymysql
import matplotlib.pyplot as plt
import numpy as np
from prettytable import PrettyTable
import pandas as pd

def generate_bar_chart():
    conn = pymysql.connect(
        host='10.114.183.17',
        user='labuser',
        password='123456',
        db='exceltodb',
        port=3306,
        charset='utf8'
    )
    cursor = conn.cursor()
    
    # 查询每个产品按test_category分类的工时
    query = """
    SELECT 
        Product, 
        test_category, 
        SUM(Test_hour * Coefficient) as category_hours
    FROM assignment
    WHERE 
        Date BETWEEN '2025-07-01' AND '2025-09-30' 
        AND product != '' 
        AND test_category != 'Operation' 
        AND special != 's'
    GROUP BY Product, test_category
    """
    cursor.execute(query)
    results = cursor.fetchall()
    conn.close()
    
    # 定义类别顺序和特殊处理规则
    category_order = ['Full ALT', 'ALT Eval', 'DVT', 'DVT Eval', 'CALT', 'CDVT', 'Carriers', 'FCNT', 'Others']
    carriers_subcategories = ['TMO', 'ATT', 'Softbank', 'Docomo']
    
    # 处理数据
    product_data = {}
    for row in results:
        product, category, hours = row
        hours = round(hours, 1)
        
        if product not in product_data:
            product_data[product] = {cat: 0 for cat in category_order}
        
        # 特殊处理Carriers类别
        if category in carriers_subcategories:
            product_data[product]['Carriers'] += hours
        # 检查是否在指定类别中
        elif category in category_order:
            product_data[product][category] += hours
        # 否则归入Others
        else:
            product_data[product]['Others'] += hours
    
    # 计算每个产品的总工时并排序
    product_totals = {product: sum(cats.values()) for product, cats in product_data.items()}
    sorted_products = sorted(product_totals.items(), key=lambda x: x[1], reverse=True)[:10]
    
    # 创建表格
    table_columns = ['产品名称'] + category_order + ['总工时']
    check_table = PrettyTable(table_columns)
    check_table.align = 'r'
    check_table.align['产品名称'] = 'l'
    
    # 准备Excel导出数据
    excel_data = []
    
    for product, total_hours in sorted_products:
        row_data = [product]
        categories = product_data[product]
        
        # 按指定顺序添加各类别工时
        for cat in category_order:
            row_data.append(categories.get(cat, 0))
        
        # 添加总工时
        row_data.append(round(total_hours, 1))
        
        check_table.add_row(row_data)
        
        # 收集Excel数据
        excel_row = {'产品名称': product, '总工时': round(total_hours, 1)}
        for cat in category_order:
            excel_row[cat] = categories.get(cat, 0)
        excel_data.append(excel_row)
    
    # 打印表格
    print("\n===== 产品工时分类汇总表 =====")
    print(check_table)
    
    # 导出Excel
    try:
        df = pd.DataFrame(excel_data)
        excel_path = '产品工时详细数据.xlsx'
        df.to_excel(excel_path, index=False)
        print(f"\n详细数据已导出至：{excel_path}")
    except Exception as e:
        print(f"\n导出Excel失败：{str(e)}")
    
    # 创建图表
    fig, ax = plt.subplots(figsize=(14, 7))
    bar_width = 0.6
    index = np.arange(len(sorted_products))
    
    # 优化颜色搭配（使用更协调的配色方案）
    colors = {
        'Full ALT': '#4287f5',     # 蓝色 - 主色调
        'ALT Eval': '#7bf542',     # 绿色
        'DVT': '#f5a442',          # 橙色
        'DVT Eval': '#f54242',     # 红色
        'CALT': '#9c42f5',         # 紫色
        'CDVT': '#f542d7',         # 粉色
        'Carriers': '#42f5e3',     # 青色
        'FCNT': '#f5f542',         # 黄色
        'Others': '#a3a3a3'        # 灰色
    }
    
    # 绘制每个产品的柱子（调整顺序为图例顺序）
    for i, (product, total_hours) in enumerate(sorted_products):
        categories = product_data[product]
        bottom = 0
        
        # 按图例顺序（从下到上）绘制类别
        for cat in reversed(category_order):
            hours = categories.get(cat, 0)
            if hours > 0:
                ax.bar(
                    i, hours, bar_width, bottom=bottom,
                    color=colors.get(cat, 'gray'),
                    label=cat if i == 0 else ""
                )
                bottom += hours
    
    # 在柱子上方标注总工时
    for i, (product, total_hours) in enumerate(sorted_products):
        ax.text(
            i, total_hours + 1,
            f'{total_hours:.1f}',
            ha='center', va='bottom', fontsize=9
        )
    
    # 设置图表属性
    ax.set_xlabel('产品名称')
    ax.set_ylabel('总工时')
    ax.set_title('2025年三季度产品工时分布（按测试类别分类）')
    ax.set_xticks(index)
    ax.set_xticklabels([p for p, _ in sorted_products], rotation=45, ha='right')
    
    # 添加图例（移至右上角，调整边框和透明度）
    handles, labels = ax.get_legend_handles_labels()
    unique = {l: h for h, l in zip(handles, labels) if l}
    ordered_handles = []
    ordered_labels = []
    for cat in category_order:
        if cat in unique:
            ordered_handles.append(unique[cat])
            ordered_labels.append(cat)
    
    ax.legend(
        ordered_handles, ordered_labels, 
        title='测试类别', 
        loc='upper right', 
        bbox_to_anchor=(0.95, 0.95),  # 调整图例位置
        frameon=True,  # 显示图例边框
        framealpha=0.8,  # 设置图例透明度
        edgecolor='gray',  # 图例边框颜色
        facecolor='white',  # 图例背景颜色
        fontsize=9  # 图例文字大小
    )
    
    # 添加网格线，提高可读性
    ax.grid(axis='y', linestyle='--', alpha=0.7)
    
    plt.tight_layout()
    plt.show()

generate_bar_chart()